import pandas as pd
import numpy as np
import pywt
import matplotlib.pyplot as plt

# 数据处理函数
def process_data():
    INPUT_FILE = "ori_message.xlsx"
    excel_file = pd.ExcelFile(INPUT_FILE)
    years = ['2016', '2017', '2018', '2019', '2020', '2021']
    all_monthly_data = []
    for year in years:
        df = excel_file.parse(year)
        df['年'] = df['年'].ffill().astype(int)
        df['月'] = df['月'].ffill().astype(int)
        grouped = df.groupby(['年', '月'])['流量(m3/s)'].mean().reset_index()
        all_monthly_data.append(grouped)
    result = pd.concat(all_monthly_data, ignore_index=True)
    result = result.sort_values(by=['年', '月']).reset_index(drop=True)
    # 生成日期列
    result['日期'] = pd.to_datetime(
        result['年'].astype(str) + '-' +
        result['月'].astype(str).str.zfill(2) + '-01'
    )
    return result[['日期', '流量(m3/s)']]

# 获取处理后的数据
df = process_data()
time_series = df['流量(m3/s)'].values  # 流量时间序列

# 小波变换参数
dt = 1  # 时间间隔（月）
scales = np.arange(1, 61)  # 调整尺度范围为 1 - 60，更贴合图 14 的横轴范围
wavelet = 'morl'  # PyWavelets 的 Morlet 小波

# 执行连续小波变换
[coefficients, _] = pywt.cwt(time_series, scales, wavelet, dt)

# 计算小波系数方差（对每个尺度的系数计算方差）
wavelet_variance = np.var(coefficients, axis=1)

# 绘制小波系数方差图
plt.figure(figsize=(10, 6))
plt.plot(scales, wavelet_variance)
plt.xlabel('时间尺度')
plt.ylabel('小波方差')
plt.title('2016 - 2021 年每月平均流量小波系数方差图')
plt.show()